program tedistboundsiter, rclass
	syntax varlist(min=4 numeric) [if] [in], ///
	[tevalues(numlist)] [y0values(numlist)] 
	tempname deltacdf cate cdeltacdf
	
	_tedistbounds `varlist',tevalues(`tevalues') y0values(`y0values')
	display "finished underscore"
	mat `deltacdf' = r(deltacdf)
	if "`tevalues'"=="" local tevalues 0

	local tind=0
	foreach t in `tevalues' {
		local ++tind
		return scalar d`tind' = `deltacdf'[`tind',1]
		return scalar dlb`tind'=`deltacdf'[`tind',2]
		return scalar dub`tind'=`deltacdf'[`tind',3]
	}
	if "`y0values'"!="" {
		mat `cate'=r(cate)
		mat `cdeltacdf' = r(cdeltacdf)
		local ytind=0
		local yind=0
		foreach y in `y0values' {
			local ++yind
			return scalar c`yind'  =`cate'[`yind',1]
			return scalar clb`yind'=`cate'[`yind',2]
			return scalar cub`yind'=`cate'[`yind',3]
			foreach t in `tevalues' {
				local ++ytind
				return scalar cdy`ytind' = `cdeltacdf'[`ytind',1]
				return scalar cdte`ytind' = `cdeltacdf'[`ytind',2]
				return scalar cdlb`ytind'=`cdeltacdf'[`ytind',3]
				return scalar cdub`ytind'=`cdeltacdf'[`ytind',4]
			}
		}
	}
	
end

program _tedistbounds, rclass
	version 13
	local version
	/* 
	This program constructs bounds on the distribution of treatment effects based on stochastic increasingness
	example syntax is:
	tedistbounds y d z s
	S must be independent of z.
	It can also test whether potential outcomes are positively correlated
	Since this program will be incorporating exchangeability, I actually want
	to be using the cdf bounds conditional on S alone (not Y(0) and S)
	to aggregate.
	Version control:
	both: imposes stochastic increasingness in both directions: Y(1) wrt Y(0) and vice versa
	Now the vanilla version automatically imposes SI both directions
	ci: does inference based on bootstrap and critical values proposed by Imbens and Manski
	*/
	syntax varlist(min=4 numeric) [if] [in], ///
	[tevalues(numlist)] [y0values(numlist)] 
	marksample touse
	* initialize temporary names, variables, and files:
	tempname beta tevaluesmat tlabelmat wjhatmat deltacdf y0valuesmat cate cdeltacdf
	tempvar kappa wj wjt y0spline fy0shat fy1tshat fy1shat fy0tshat ones zeros requestedvalue requestedtvalue
	tokenize `varlist'
	local y "`1'"
	macro shift
	local d "`1'"
	macro shift
	local z "`1'"
	macro shift
	local s "`*'"
	macro shift
	
	* make dataset of requested y0vals
	preserve
	clear	
	if "`y0values'"!="" {
		matrix `y0valuesmat'=J(1,1,.)
		foreach y0 in `y0values' {
			matrix `y0valuesmat'=`y0valuesmat' \ `y0'
		}
		qui svmat `y0valuesmat'
		rename `y0valuesmat'1 `y'
		qui drop if `y'==.
		gen byte `requestedvalue'=1
	}
	else {
		qui gen `y'=.
		qui gen `requestedvalue'=.
	}
	tempfile requestedy0vals y0vals y1vals
	qui save `requestedy0vals'			
	restore
	preserve
	qui keep if `touse'
	keep `y' `d' `z' `s'
	* first calculate Abadie's kappa
	* decide if treatment is endogenous
	qui count if `d'!=`z'
	local endog=r(N)>0
	if `endog'==0 {
		gen `kappa'=1
	}
	else {
		* estimate tau=E[Z]
		qui sum `z'
		local tau=r(mean)
		* now construct Abadie's kappa weights:
		qui gen `kappa'=1-`d'*(1-`z')/(1-`tau')-(1-`d')*`z'/`tau'
	}
	gen `ones'=1
	gen `zeros'=0
	* generate interactions of s with d
	local sxd
	local sx0
	foreach svar of varlist `s' {
		tempvar `svar'`d' `svar'0
		qui gen ``svar'`d''=`d'*`svar'
		qui gen ``svar'0'=0
		local sxd `sxd' ``svar'`d''
		local sx0 `sx0' ``svar'0'
	}

	tempfile goback
	qui save `goback'
	
	if "`y0values'"!= "" {
		* shuffle in the requested values
		append using `requestedy0vals'
		* generate Y(0) spline for later integration over S:
		mkspline `y0spline'=`y' if `d'==0 | `requestedvalue'==1,cubic nknots(5)
		qui save `goback',replace
		qui keep if `requestedvalue'==1
		keep `y' `requestedvalue' `y0spline'*
		qui save `requestedy0vals',replace
		use `goback'
		qui drop if `requestedvalue'==1
		drop `requestedvalue'
		qui save `goback',replace
	}
	foreach dval in 0 1 {
		qui keep if `d'==`dval'
		keep `y'
		rename `y' `y'`dval'
		
		qui duplicates drop
		local numy`dval'vals=_N
		sort `y'`dval'
		mkmat *
		use `goback',clear
	}
	
	mata `tevaluesmat'=J(0,1,0)
	if "`tevalues'"=="" local tevalues 0
	foreach te in `tevalues' {
		mata `tevaluesmat'=`tevaluesmat' \ `te'
	}
	
	* only make a full grid of TE values if we are computing CATE:
	if "`y0values'"!="" {
		* decide support of the treatment effect distribution.
		foreach dval in 0 1 {
			qui sum `y' if `d'==`dval'
			local max`dval'=r(max)
			local min`dval'=r(min)
		}
		local t1=`max1'-`min0'
		local t0=`min1'-`max0'
		local gridn=25
		local tstep=(`t1'-`t0')/`gridn'
		mata `tlabelmat'=1::`gridn'+1
		mata `tlabelmat'=(`t0':+`tstep':*(`tlabelmat':-1))
	}
	else mata `tlabelmat'=J(0,1,0)
	* now shuffle in the requested evaluation points for the TE cdf:
	mata `tlabelmat'=uniqrows(floatround(`tlabelmat') \ floatround(`tevaluesmat'))
	mata `tlabelmat'=(1::length(`tlabelmat')),`tlabelmat'
	mata st_matrix("`tlabelmat'",`tlabelmat')
	local tlabelmax=rowsof(`tlabelmat')
	clear
	qui svmat `tlabelmat'
	rename `tlabelmat'1 tlabel
	rename `tlabelmat'2 tval
	
	tempfile tlabels
	qui save `tlabels'
	use `goback',clear
	* initialize variables to hold estimated cdfs:
	qui gen `fy0shat'=.
	qui gen `fy1tshat'=.
	qui gen `fy1shat'=.
	qui gen `fy0tshat'=.
	forvalues tlabel=1/`tlabelmax' {
		qui gen deltayslb`tlabel'=.
		qui gen deltaysub`tlabel'=.
		qui gen deltaysaltlb`tlabel'=.
		qui gen deltaysaltub`tlabel'=.
	}
	
		
	* first calculate CDF of Y(0) at Y(1)-t conditional on Y(1) and S, for later comparison. 
	* Don't forget strict inequality, and minus t
	forvalues j=1/`numy1vals' {
		local y1val=`y'1[`j',1]
		qui gen `wj'=float(`y')<=float(`y1val')
		mata output=wls("`wj' `kappa' `d' `s' `sxd'","`d' `s' `sxd'")
		mata st_matrix("`beta'",*output[1])
		mata st_matrix("`wjhatmat'",*output[3])
		svmat `wjhatmat'
		qui replace `fy1shat'=max(0,min(1,`wjhatmat')) if `d'==1 & float(`y')==float(`y1val')
		drop `wjhatmat'
		* now get fy0 at y1-t
		forvalues tlabel=1/`tlabelmax' {
			local t=`tlabelmat'[`tlabel',2]
			qui gen `wjt'=float(`y')<float(`y1val'-`t')
			
			* do kappa-weighted regression. Need to remember whats up with the fitted values--why evaluate as if d==0?
			* it's because I want the expectation of Y(0), but I'm evaluating it at each Y(1) value
			mata output=wls("`wjt' `kappa' `s' `sxd' `d' ","`s' `sx0' `zeros' ")
			mata st_matrix("`beta'",*output[1])
			* now construct fitted values:
			mata st_matrix("`wjhatmat'",*output[3])
			svmat `wjhatmat'
			qui replace `fy0tshat'=max(0,min(1,`wjhatmat')) if `d'==1 & float(`y')==float(`y1val')
			drop `wjt' `wjhatmat'
			qui replace deltaysaltlb`tlabel'=max((`fy0tshat'-`fy1shat')/(1-`fy1shat'),0) if `d'==1 & float(`y')==float(`y1val')
			qui replace deltaysaltub`tlabel'=min(`fy0tshat'/`fy1shat',1) if `d'==1 & float(`y')==float(`y1val')
			qui replace deltaysaltlb`tlabel'= 1 if `d'==1 & float(`y')==float(`y1val') & float(`fy1shat')==float(1) & float(`fy0tshat')==float(1)
			qui replace deltaysaltub`tlabel'= 0 if `d'==1 & float(`y')==float(`y1val') & float(`fy0tshat')==float(0) 
			
		}
		drop `wj' 
		
	}
	
	* Now integrate over y(1) to get bounds by S only
	* Why? because conditional on S we can pick the tighter of the bounds in either direction
	forvalues tlabel=1/`tlabelmax' {
		local t=`tlabelmat'[`tlabel',2]
		foreach b in lb ub {
			* now integrate out Y(1) but leave s for imposing the other direction
			unab vlist: deltaysalt`b'`tlabel' `kappa' `s' 
			unab pvlist: `s' 
			mata output=wls("`vlist'","`pvlist'")
			tempname deltasalt`b'`tlabel'
			mata st_matrix("`deltasalt`b'`tlabel''",*output[3])
			svmat `deltasalt`b'`tlabel''
			rename `deltasalt`b'`tlabel''1 deltasalt`b'`tlabel'
			
			* finally, do one minus them to compare to other direction
			qui replace deltasalt`b'`tlabel'=1-deltasalt`b'`tlabel'
			
		}
	}
	* now do it imposing Y(1) SI wrt Y(0) (given S)
	
	forvalues j=1/`numy0vals' {
		*display "on z-step `zlabel' of `numzlabels', j = `j' of `n0'"
		local y0val=`y'0[`j',1]
		qui gen `wj'=float(`y')<=float(`y0val')
		
		mata output=wls("`wj' `kappa' `d' `s' `sxd'","`d' `s' `sxd'")
		mata st_matrix("`beta'",*output[1])
		mata st_matrix("`wjhatmat'",*output[3])
		svmat `wjhatmat'
		qui replace `fy0shat'=max(0,min(1,`wjhatmat')) if `d'==0 & float(`y')==float(`y0val')
		drop `wjhatmat'
		* now get fy1 at y0+t
		forvalues tlabel=1/`tlabelmax' {
			local t=`tlabelmat'[`tlabel',2]
			qui gen `wjt'=float(`y')<=float(`y0val'+`t')
			* do kappa-weighted regression. Need to remember whats up with the fitted values--why evaluate as if d==1?
			mata output=wls("`wjt' `kappa' `s' `sxd' `d' ","`s' `s' `ones' ")
			mata st_matrix("`beta'",*output[1])
			* now construct fitted values:
			mata st_matrix("`wjhatmat'",*output[3])
			svmat `wjhatmat'
			qui replace `fy1tshat'=max(0,min(1,`wjhatmat')) if `d'==0 & float(`y')==float(`y0val')
			drop `wjt' `wjhatmat'
			qui replace deltayslb`tlabel'=max((`fy1tshat'-`fy0shat')/(1-`fy0shat'),0) if `d'==0 & float(`y')==float(`y0val')
			qui replace deltaysub`tlabel'=min(`fy1tshat'/`fy0shat',1) if `d'==0 & float(`y')==float(`y'0[`j',1])
			qui replace deltayslb`tlabel'= 1 if `d'==0 & float(`y')==float(`y0val') & float(`fy0shat')==float(1) & float(`fy1tshat')==float(1)
			qui replace deltaysub`tlabel'= 0 if `d'==0 & float(`y')==float(`y0val') & float(`fy1tshat')==float(0) 
			*display "j = `j', y0 = `y0val', t = `t', fy0=`fy0shat', fy1 = `fy1tshat'"
			*sum `fy0shat' `fy1tshat' `delta'yslb`tlabel' `delta'ysub`tlabel' if `d'==0 & float(`y')==float(`y'0[`j',1])
		}
		drop `wj' 
		
	}
	* now integrate over s to get bounds by Y(0) only
	* and also over y(0) to get bounds by S only (if requested
	mata cdf=J(0,3,.)
	
	
	* now shuffle in requested Y(0) values, if any: 
	if "`y0values'"!="" {
		append using `requestedy0vals'
	}
	
	forvalues tlabel=1/`tlabelmax' {
		local t=`tlabelmat'[`tlabel',2]
		foreach b in lb ub {
			if "`y0values'"!="" {
				unab vlist: deltays`b'`tlabel' `kappa' `y0spline'* 
				unab pvlist: `y0spline'* 
				mata output=wls("`vlist'","`pvlist'")
				tempname delta`y'`b'`tlabel'
				mata st_matrix("`delta`y'`b'`tlabel''",*output[3])
				svmat `delta`y'`b'`tlabel''
				rename `delta`y'`b'`tlabel''1 delta`y'`b'`tlabel'
				qui replace delta`y'`b'`tlabel'=min(1,max(0,delta`y'`b'`tlabel'))
				*twoway scatter `delta'`y'`b'`tlabel' `delta'ys`b'`tlabel' `y',title("`b': t = `t'")
				*pause
				*local toagg `delta'ys`b'`tlabel'
			}
			* now integrate out Y(0) but leave s for imposing the other direction and (if desired) exchangeability
			unab vlist: deltays`b'`tlabel' `kappa' `s' 
			unab pvlist: `s' 
			mata output=wls("`vlist'","`pvlist'")
			tempname deltas`b'`tlabel'
			mata st_matrix("`deltas`b'`tlabel''",*output[3])
			svmat `deltas`b'`tlabel''
			rename `deltas`b'`tlabel''1 deltas`b'`tlabel'
			qui replace deltas`b'`tlabel'=min(1,max(0,deltas`b'`tlabel'))
			* now take tighter of two bounds between SI imposed in either direction. Remember, switch ub and lb
			if "`b'"=="ub" qui replace deltas`b'`tlabel'=min(deltas`b'`tlabel',deltasaltlb`tlabel')
			else qui replace deltas`b'`tlabel'=max(deltas`b'`tlabel',deltasaltub`tlabel')
			local toagg deltas`b'`tlabel'
			* now average to get overall cdf
			mata output=wls("`toagg' `kappa'","")
			mata cdf`b'=*output[1]
		}
		mata cdf=cdf\(`t',min((1,max((0,cdflb)))),min((1,max((0,cdfub)))))
	}
	mata cdfrequested=J(0,3,.)
	foreach t in `tevalues' {
		mata cdfrequested=cdfrequested\select(cdf,floatround(`t'):==floatround(cdf[.,1]))
	} 
	mata cdfrequested
	mata st_matrix("`deltacdf'",cdfrequested)
	mat colnames `deltacdf'=tevalue cdf_lb cdf_ub
	
	
		
	
	if "`y0values'"!="" {
		tempfile deltays
		qui keep if `d'==0 | `requestedvalue'==1
		keep delta`y'lb* delta`y'ub* `y'
		qui duplicates drop
		rename `y' `y'0val
		qui reshape long delta`y'lb delta`y'ub, i(`y'0val) j(tlabel)
		qui merge m:1 tlabel using `tlabels'
		drop _merge
		drop tlabel
		
		order *, alphabetic
		order `y'0val tval , first
		sort `y'0val tval
		qui save `deltays'
		keep tval
		tempfile tvals
		qui save `tvals'
		use `deltays'
		keep `y'0val delta`y'lb
		qui replace delta`y'lb=min(max(0,delta`y'lb),1)
		sort `y'0val delta`y'lb
		qui merge 1:1 _n using `tvals'
		drop _merge
		tempfile lbs
		qui save `lbs'
		use `deltays'
		keep `y'0val delta`y'ub
		qui replace delta`y'ub=min(max(0,delta`y'ub),1)
		sort `y'0val delta`y'ub
		qui merge 1:1 _n using `tvals'
		drop _merge
		*li `y'0val tval
		qui merge 1:1 `y'0val tval using `lbs'
		drop _merge
	
		rename tval deltaval
		sort deltaval `y'0val
		order deltaval `y'0val  delta`y'lb delta`y'ub
		rename (delta`y'lb delta`y'ub) (delta`y'cdflb delta`y'cdfub)
		
		keep deltaval `y'0val delta`y'cdflb delta`y'cdfub
		
		qui keep if deltaval !=. & `y'0val !=.
		sort `y'0val deltaval
		egen tvar=group(deltaval)
		qui tsset `y'0val tvar
		sort `y'0val tvar
		qui gen densylb=d.delta`y'cdflb
		qui gen densyub=d.delta`y'cdfub
		by `y'0val: egen ate`y'lb=total(densyub*deltaval)
		by `y'0val: egen ate`y'ub=total(densylb*deltaval)
		egen tag=tag(`y'0val)
		qui replace ate`y'lb=. if tag!=1
		qui replace ate`y'ub=. if tag!=1
		qui save `deltays',replace
		keep `y'0val deltaval delta`y'cdflb delta`y'cdfub
		rename `y'0val `y'
		qui merge m:1 `y' using `requestedy0vals'
		drop _merge
		qui keep if `requestedvalue'==1
		keep `y' deltaval delta`y'cdflb delta`y'cdfub
		
		* now need to keep only requested tvalues
		qui gen `requestedtvalue'=0
		foreach t in `tevalues' {
			qui replace `requestedtvalue'=1 if float(deltaval)==float(`t')
		}
		qui keep if `requestedtvalue'==1
		mkmat `y' deltaval delta`y'cdflb delta`y'cdfub,matrix(`cdeltacdf')
		use `deltays',clear
		keep `y'0val ate`y'lb ate`y'ub
		rename `y'0val `y'
		qui keep if `y'!=. & ate`y'lb!=. & ate`y'ub!=.
		qui duplicates drop
		qui merge 1:1 `y' using `requestedy0vals'
		drop _merge
		qui keep if `requestedvalue'==1
		keep `y' ate`y'lb ate`y'ub
		rename (ate`y'lb ate`y'ub) (cate_lb cate_ub)
		mkmat `y' cate_lb cate_ub,matrix(`cate')
	}
	restore
	
	return matrix deltacdf = `deltacdf'
	if "`y0values'"!="" {
		return matrix cate=`cate'
		return matrix cdeltacdf = `cdeltacdf'
	}
end

* mata programs
mata:
	pointer vector wls(string scalar varlist,string scalar pvarlist)
	{ // does weighted least squares of depvar on x1 x2 x3 with weights kappa.
	//feed it: depvar kappa x1 x2 x3 ...
	// pvarlist is a list of variables that will be used to construct the predicted values
	
	vars=tokens(varlist)
	pvars=tokens(pvarlist)
	y=st_data(.,vars[1])
	kappa=st_data(.,vars[2])
	if (length(vars)>=3) {
		xs=st_data(.,vars[3..length(vars)])
		xs=xs,J(rows(xs),1,1)
	}
	else xs=J(rows(y),1,1)
	if (length(pvars)>=1) {
		pxs=st_data(.,pvars[1..length(pvars)])
		pxs=pxs,J(rows(pxs),1,1)
	}
	else pxs=J(rows(y),1,1)
	// set missing rows in y to be zero in kappa
	kappa=(y:!=.):*kappa
	xkx=cross(xs,kappa,xs)
	xky=cross(xs,kappa,y)
	xkxinv=invsym(xkx)
	beta=xkxinv*xky
	u=y-xs*beta
	u2k=u:^2:*kappa
	varcov=xkxinv*cross(xs,u2k,xs)*xkxinv
	yhat=pxs*beta
	return((&beta,&varcov,&yhat))
}
end


